// prefetcher.cc (Version 3)
#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <sstream>
#include <csignal>
#include <ctime>
#include <unistd.h>
#include <sys/syscall.h>
#include <fcntl.h>
#include <cerrno>
#include <cstring>
#include <unordered_map>
#include <iomanip>
#include <deque>    // V3: For storing history
#include <utility>  // V3: For std::pair

extern "C" {
#include <bpf/libbpf.h>
#include <bpf/bpf.h>
#include "prefetcher.skel.h"
#include "prefetcher.h"
}

// ============================================================================
// V3: 用户态宏观预测模型核心
// ============================================================================

// --- 模型参数 ---
#define HISTORY_SIZE 5              // 记录最近多少次缺页事件用于分析
#define ANALYSIS_THRESHOLD 3        // 历史记录达到多少条时开始分析
#define CONFIDENCE_THRESHOLD 3      // 连续多少次分析结果一致则认为模式稳定
#define PREFETCH_HORIZON_S 0.5      // 预测未来多少秒的数据量
#define RATE_TOLERANCE 0.20         // 速率变化容忍度 (20%)

// V3: 每次缺页事件的详细记录 (在文件偏移量空间)
struct FaultRecord {
    struct timespec timestamp;
    off_t file_offset;
};

// V3: 追踪每个线程的访问流状态 (基于速率和方向)
struct StreamTracker {
    uint32_t pid;
    uint32_t tid;
    std::deque<FaultRecord> history;

    // 宏观状态
    int direction;                  // 1: 正向, -1: 反向
    double rate_bytes_per_sec;
    int confidence;
    std::string file_path;          // 追踪的文件路径

    StreamTracker(uint32_t p, uint32_t t, const std::string& path)
        : pid(p), tid(t), direction(0), rate_bytes_per_sec(0.0), confidence(0), file_path(path) {}
};

// 全局状态：使用TID作为key，追踪所有检测到的访问流
static std::unordered_map<uint32_t, StreamTracker> stream_trackers;

// ============================================================================
// 全局变量和辅助函数
// ============================================================================

std::ofstream log_file;
static volatile bool exiting = false;

// 统计计数器
static long faults_processed_count = 0;
static long faults_in_files_count = 0; // V3 新增: 属于文件映射的缺页
static long pattern_detected_count = 0;
static long readahead_success_count = 0;
static long readahead_failure_count = 0;

// V3: 时间戳辅助函数
double get_elapsed_seconds(const struct timespec& start, const struct timespec& end) {
    return (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9;
}

std::string get_timestamp() {
    time_t now = time(0);
    struct tm tstruct;
    char buf[80];
    tstruct = *localtime(&now);
    strftime(buf, sizeof(buf), "%Y-%m-%d %X", &tstruct);
    return buf;
}

void log_message(const std::string& msg) {
    std::string full_msg = "[" + get_timestamp() + "] " + msg;
    std::cout << full_msg << std::endl;
    if (log_file.is_open()) {
        log_file << full_msg << std::endl;
    }
}

// 地址转换函数 (保持不变)
bool find_file_offset(uint32_t pid, unsigned long vma_start, int& fd, off_t& file_offset, std::string& file_path) {
    std::string maps_path = "/proc/" + std::to_string(pid) + "/maps";
    std::ifstream maps_file(maps_path);
    if (!maps_file.is_open()) return false;
    std::string line;
    while (std::getline(maps_file, line)) {
        std::stringstream ss(line);
        std::string address_range, perms, offset_str, dev, inode, pathname;
        ss >> address_range >> perms >> offset_str >> dev >> inode >> pathname;
        if (pathname.empty() || pathname[0] == '[') continue;
        
        unsigned long start, end, file_map_offset;
        sscanf(address_range.c_str(), "%lx-%lx", &start, &end);
        sscanf(offset_str.c_str(), "%lx", &file_map_offset);

        if (vma_start >= start && vma_start < end) {
            fd = open(pathname.c_str(), O_RDONLY);
            if (fd < 0) return false;
            file_offset = (vma_start - start) + file_map_offset;
            file_path = pathname;
            return true;
        }
    }
    return false;
}

// Ring buffer 事件处理函数 (V3: 核心逻辑)
int handle_event(void *ctx, void *data, size_t data_sz) {
    const struct fault_info *e = (struct fault_info*)data;
    faults_processed_count++;

    // 1. 数据转换与记录
    struct timespec ts_now;
    clock_gettime(CLOCK_MONOTONIC, &ts_now);

    int fd = -1;
    off_t current_offset;
    std::string file_path;
    if (!find_file_offset(e->pid, e->addr, fd, current_offset, file_path)) {
        // 这不是文件映射的缺页（例如，匿名内存），忽略
        if (fd != -1) close(fd);
        return 0;
    }
    faults_in_files_count++;

    // 2. 查找或创建追踪器
    auto it = stream_trackers.find(e->tid);
    if (it == stream_trackers.end() || it->second.file_path != file_path) {
        // 如果是新TID，或者同一个TID开始访问一个新文件，则创建新的追踪器
        if (it != stream_trackers.end()) {
             stream_trackers.erase(it);
        }
        it = stream_trackers.emplace(e->tid, StreamTracker(e->pid, e->tid, file_path)).first;
        log_message("New stream for TID " + std::to_string(e->tid) + " on file: " + file_path);
    }
    
    // 3. 更新历史记录
    StreamTracker& tracker = it->second;
    tracker.history.push_back({ts_now, current_offset});
    if (tracker.history.size() > HISTORY_SIZE) {
        tracker.history.pop_front();
    }
    
    // 4. 宏观模式分析
    if (tracker.history.size() < ANALYSIS_THRESHOLD) {
        if (fd != -1) close(fd);
        return 0; // 历史数据不足，无法分析
    }

    const auto& first_record = tracker.history.front();
    const auto& last_record = tracker.history.back();
    
    double delta_time = get_elapsed_seconds(first_record.timestamp, last_record.timestamp);
    if (delta_time < 1e-6) { // 时间太短，避免除零
        if (fd != -1) close(fd);
        return 0;
    }
    
    off_t delta_offset = last_record.file_offset - first_record.file_offset;
    double current_rate = delta_offset / delta_time;
    int current_direction = (delta_offset > 0) ? 1 : ((delta_offset < 0) ? -1 : 0);

    std::stringstream log_msg;
    log_msg << "TID " << e->tid << ": Analyzing window. "
            << "Offset delta: " << delta_offset / 1024.0 << " KB, "
            << "Time delta: " << std::fixed << std::setprecision(3) << delta_time << " s, "
            << "Rate: " << std::fixed << std::setprecision(2) << current_rate / (1024*1024) << " MB/s";
    log_message(log_msg.str());

    // 5. 更新置信度
    bool is_consistent = true;
    if (current_direction == 0) {
        is_consistent = false; // 原地踏步不是稳定模式
    }
    if (current_direction != tracker.direction) {
        is_consistent = false;
    }
    if (tracker.rate_bytes_per_sec > 1.0) { // 如果已有基准速率
        double rate_diff = std::abs(current_rate - tracker.rate_bytes_per_sec) / tracker.rate_bytes_per_sec;
        if (rate_diff > RATE_TOLERANCE) {
            is_consistent = false;
        }
    }

    if (is_consistent) {
        tracker.confidence++;
        // 平滑更新速率，使其更稳定
        tracker.rate_bytes_per_sec = (tracker.rate_bytes_per_sec * 0.7) + (current_rate * 0.3);
    } else {
        tracker.confidence = 1;
        tracker.direction = current_direction;
        tracker.rate_bytes_per_sec = current_rate;
    }
    log_message("  -> Confidence: " + std::to_string(tracker.confidence) + ", Direction: " + std::to_string(tracker.direction));

    // 6. 预测与执行
    if (tracker.confidence >= CONFIDENCE_THRESHOLD && tracker.direction == 1) { // 只对正向读取预取
        pattern_detected_count++;
        
        size_t prefetch_bytes = tracker.rate_bytes_per_sec * PREFETCH_HORIZON_S;
        // 对预取大小做一些合理性约束
        prefetch_bytes = std::max((size_t) (128 * 1024), std::min(prefetch_bytes, (size_t) (8 * 1024 * 1024)));

        off_t prefetch_offset = last_record.file_offset;

        log_msg.str(""); log_msg.clear();
        log_msg << "  -> STEADY PATTERN! Prefetching " << prefetch_bytes / 1024 << " KB "
                << "from offset " << prefetch_offset;
        log_message(log_msg.str());
        
        long ret = syscall(SYS_readahead, fd, prefetch_offset, prefetch_bytes);
        if (ret == 0) {
            readahead_success_count++;
            log_message("  -> readahead() SUCCEEDED.");
        } else {
            readahead_failure_count++;
            log_message("  -> readahead() FAILED. Errno: " + std::to_string(errno));
        }

        // 成功预取后，清空历史，开始新的观测周期
        tracker.history.clear();
        tracker.confidence = 0;
    }

    if (fd != -1) close(fd);
    return 0;
}

// libbpf 日志回调函数
int print_libbpf_log(enum libbpf_print_level level, const char *format, va_list args) {
    if (level == LIBBPF_WARN) {
        char buffer[1024];
        vsnprintf(buffer, sizeof(buffer), format, args);
        log_message("libbpf: " + std::string(buffer));
    }
    return 0;
}

static void sig_handler(int sig) { exiting = true; }

// ============================================================================
// Main 函数
// ============================================================================
int main(int argc, char **argv) {
    if (argc > 1 && (std::string(argv[1]) == "-h" || std::string(argv[1]) == "--help")) {
        std::cout << "Usage: " << argv[0] << " [target_pid]" << std::endl;
        return 0;
    }

    log_file.open("prefetcher.log", std::ios::out | std::ios::trunc);
    log_message("Prefetcher V3 starting...");

    libbpf_set_print(print_libbpf_log);

    struct prefetcher_bpf *skel = prefetcher_bpf__open_and_load();
    if (!skel) {
        log_message("Error: Failed to open and load BPF skeleton");
        return 1;
    }

    struct config cfg = {};
    if (argc > 1) cfg.target_pid = std::stoi(argv[1]);
    const int map_key = 0; 
    bpf_map_update_elem(bpf_map__fd(skel->maps.config_map), &map_key, &cfg, BPF_ANY);
    log_message(argc > 1 ? "Targeting PID: " + std::string(argv[1]) : "Monitoring all processes.");

    int err = prefetcher_bpf__attach(skel);
    if (err) {
        log_message("Error: Failed to attach BPF skeleton: " + std::string(strerror(-err)));
        prefetcher_bpf__destroy(skel);
        return 1;
    }

    log_message("eBPF program loaded and attached successfully.");
    signal(SIGINT, sig_handler);
    signal(SIGTERM, sig_handler);

    struct ring_buffer *rb = ring_buffer__new(bpf_map__fd(skel->maps.events), handle_event, NULL, NULL);
    if (!rb) {
        err = -1;
        log_message("Error: Failed to create ring buffer");
        goto cleanup;
    }

    log_message("Waiting for events... Press Ctrl+C to exit.");
    while (!exiting) {
        err = ring_buffer__poll(rb, 100);
        if (err == -EINTR) { err = 0; break; }
        if (err < 0) { log_message("Error polling ring buffer: " + std::to_string(err)); break; }
    }

cleanup:
    ring_buffer__free(rb);
    prefetcher_bpf__destroy(skel);
    
    log_message("Prefetcher stopped. Final statistics:");
    std::stringstream final_stats;
    final_stats << "  - Total page faults processed: " << faults_processed_count << std::endl;
    final_stats << "  - Page faults on file mappings: " << faults_in_files_count << std::endl;
    final_stats << "  - Steady patterns detected: " << pattern_detected_count << std::endl;
    final_stats << "  - Successful readahead calls: " << readahead_success_count << std::endl;
    final_stats << "  - Failed readahead calls: " << readahead_failure_count;
    log_message(final_stats.str());

    if (log_file.is_open()) log_file.close();
    return err < 0 ? -err : 0;
}